from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

from llm.glm_llm import glm_lite, glm

template = """Given the user question below, classify it as either being about `company` or `law`.
    
Do not respond with more than one word.

<question>
{question}
</question>

Classification:"""

router_chain = PromptTemplate.from_template(template) | glm | StrOutputParser()


if __name__ == "__main__":
    from utils import read_jsonl
    from tqdm import tqdm
    import jsonlines

    question_file = "../data/question.jsonl"
    queries = read_jsonl(question_file)
    result_file = "../data/topic.jsonl"

    for query in tqdm(queries):
        query["topic"] = router_chain.invoke(query)

    with jsonlines.open(result_file, "w") as json_file:
        for query in queries:
            json_file.write(query)
